import os
import re
import asyncio
import random
import json
from pathlib import Path
from urllib.parse import urlencode
from playwright.async_api import async_playwright


def read_env(key, default=''):
    v = os.getenv(key)
    if v:
        return v
    env_path = Path(__file__).resolve().parents[1] / 'luogu-snap' / '.env'
    if env_path.exists():
        for line in env_path.read_text(encoding='utf-8').splitlines():
            if line.strip().startswith(key + '='):
                return line.split('=', 1)[1].strip()
    return default


def sanitize(s):
    return (s or '').replace('\n', ' ').replace('\r', ' ').strip()


async def collect_links(page, user_id, page_no):
    url = f'https://www.luogu.com.cn/record/list?user={user_id}&page={page_no}'
    await page.goto(url)
    await page.wait_for_timeout(800)
    links = await page.evaluate("""
    (() => {
      const as = Array.from(document.querySelectorAll('a'));
      return Array.from(new Set(as.map(a => a.href).filter(h => h.includes('/record/') && !h.includes('/record/list'))));
    })()
    """)
    return links


async def fetch_problem_content(page, pid):
    if not pid:
        return ''
    url = f'https://www.luogu.com.cn/problem/{pid}'
    try:
        await page.goto(url)
        await page.wait_for_timeout(800)
        content = await page.evaluate(r"""
        (() => {
          const wanted = ['题目描述','问题描述','输入格式','输出格式','样例输入','样例输出','样例解释','提示','说明'];
          const headings = Array.from(document.querySelectorAll('h2, h3, .lfe-h2, .lfe-h3'));
          const sections = [];
          for (const h of headings) {
            const ht = h.textContent.trim();
            if (wanted.some(w => ht.includes(w))) {
              let seg = ht + '\n';
              let cur = h.nextElementSibling;
              while (cur && !headings.includes(cur)) {
                const tn = cur.tagName ? cur.tagName.toLowerCase() : '';
                if (['p','div','section','pre','code','ul','ol','li'].includes(tn)) {
                  seg += cur.textContent.trim() + '\n';
                }
                cur = cur.nextElementSibling;
              }
              sections.push(seg.trim());
            }
          }
          if (!sections.length) {
            const main = document.querySelector('main') || document.body;
            return main ? main.textContent : '';
          }
          return sections.join('\n\n');
        })()
        """)
        return sanitize(content)[:8000]
    except Exception:
        return ''


async def extract_detail(page, url):
    await page.goto(url)
    await page.wait_for_timeout(800)
    data = await page.evaluate(r"""
    (() => {
      const tEl = document.querySelector('[class*="title"], h1, .lg-title');
      const title = tEl ? tEl.textContent : '';
      const metaText = Array.from(document.querySelectorAll('div, span, p')).map(e => e.textContent).join('\n');
      const codeEl = document.querySelector('pre code') || document.querySelector('pre') || document.querySelector('code');
      const code = codeEl ? codeEl.textContent : '';
      const pm = metaText.match(/\b([A-Z]{1,3}\d{1,6})\b/);
      const problemId = pm ? pm[1] : '';
      const rm = location.pathname.match(/record\/(\d+)/);
      const rid = rm ? rm[1] : '';
      let lang = '';
      const lm1 = metaText.match(/编程语言[:：]\s*([A-Za-z0-9\+\#]+)/);
      const lm2 = metaText.match(/语言[:：]\s*([A-Za-z0-9\+\#]+)/);
      if (lm1) lang = lm1[1]; else if (lm2) lang = lm2[1];
      let submitTime = '';
      const timeLabel = Array.from(document.querySelectorAll('span')).find(s => (s.textContent || '').includes('提交时间'));
      if (timeLabel) {
        // 尝试读取同一行的第二个 span 作为时间值
        const p = timeLabel.parentElement;
        if (p) {
          const spans = Array.from(p.querySelectorAll('span'));
          if (spans.length >= 2) submitTime = (spans[1].textContent || '').trim();
        }
        // 兜底：尝试直接读取兄弟节点文本
        if (!submitTime && timeLabel.nextElementSibling) submitTime = (timeLabel.nextElementSibling.textContent || '').trim();
      }
      if (!submitTime) {
        const tm = metaText.match(/提交时间[:：]\s*([^\n]+)/);
        submitTime = tm ? tm[1].trim() : '';
      }
      // 结果与统计（逐测试点聚合）
      const cases = Array.from(document.querySelectorAll('.test-case'));
      let totalScore = 0;
      let allAC = true;
      let maxMs = 0;
      let maxKB = 0;
      for (const c of cases) {
        const statusEl = c.querySelector('.status');
        const s = statusEl ? statusEl.textContent.trim() : '';
        if (s !== 'AC') allAC = false;
        const infoEl = c.querySelector('.info');
        if (infoEl) {
          const mm = infoEl.textContent.match(/([0-9.]+)ms\/(\d+\.?\d*)KB/);
          if (mm) {
            const ms = parseFloat(mm[1]);
            const kb = parseFloat(mm[2]);
            if (!Number.isNaN(ms)) maxMs = Math.max(maxMs, ms);
            if (!Number.isNaN(kb)) maxKB = Math.max(maxKB, kb);
          }
        }
        const msgEl = c.querySelector('.message');
        if (msgEl) {
          const sm = msgEl.textContent.match(/得分\s*(\d+)/);
          if (sm) totalScore += parseInt(sm[1], 10);
        }
      }
      const status = cases.length ? (allAC ? 'AC' : 'Unaccepted') : (document.querySelector('.status')?.textContent.trim() || '');
      const runtime = maxMs ? (maxMs + 'ms') : '';
      const memory = maxKB ? (maxKB + 'KB') : '';
      const score = totalScore ? String(totalScore) : (metaText.match(/得分\s*(\d+)/)?.[1] || '');
      return { title, code, problemId, rid, lang, submitTime, status, score, runtime, memory };
    })()
    """)
    try:
        tab = page.locator('text=源代码').first
        await tab.click()
        await page.wait_for_timeout(500)
        el = await page.query_selector('pre code, pre, code')
        if el:
            data['code'] = await el.text_content()
    except Exception:
        pass
    if not data.get('code'):
        try:
            btn = await page.query_selector('text=查看代码, text=显示代码, text=展开')
            if btn:
                await btn.click()
                await page.wait_for_timeout(500)
                el = await page.query_selector('pre code, pre, code')
                if el:
                    data['code'] = await el.text_content()
        except Exception:
            pass
    data['title'] = sanitize(data.get('title', ''))
    st = data.get('submitTime', '')
    m = re.search(r'(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2})', st)
    if m:
        data['submitTime'] = m.group(1)
    else:
        data['submitTime'] = st.replace('提交时间', '').replace('：', '').strip()
    try:
        data['problemContent'] = await fetch_problem_content(page, data.get('problemId'))
    except Exception:
        data['problemContent'] = ''
    return data


async def render_and_shot(page, template_path, info, out_path):
    params = urlencode({
        'problem': info.get('problemId', ''),
        'title': info.get('title', ''),
        'rid': info.get('rid', ''),
        'lang': info.get('lang', ''),
        'time': info.get('submitTime', ''),
        'code': info.get('code', ''),
        'pcontent': info.get('problemContent', ''),
        'status': info.get('status', ''),
        'score': info.get('score', ''),
        'runtime': info.get('runtime', ''),
        'memory': info.get('memory', '')
    })
    url = f'file://{template_path.as_posix()}?{params}'
    await page.goto(url)
    if info.get('problemContent'):
        await page.evaluate("(t)=>{document.getElementById('problemContent').textContent=t}", info['problemContent'])
        await page.wait_for_timeout(200)
    # 动态分栏填充代码，避免长代码字体过小
    code = info.get('code', '')
    if code:
        lines = code.split('\n')
        per_col = 30
        cols = max(1, min(4, (len(lines) + per_col - 1) // per_col))
        chunks = []
        for i in range(cols):
            chunk = '\n'.join(lines[i*per_col:(i+1)*per_col])
            chunks.append(chunk)
        await page.evaluate("(n)=>{const g=document.getElementById('codeGrid');g.classList.remove('col-1','col-2','col-3','col-4');g.classList.add('col-'+n)}", cols)
        await page.evaluate("(t)=>{document.getElementById('codeCol1').textContent=t}", chunks[0] if len(chunks)>0 else '')
        await page.evaluate("(t)=>{document.getElementById('codeCol2').textContent=t}", chunks[1] if len(chunks)>1 else '')
        await page.evaluate("(t)=>{document.getElementById('codeCol3').textContent=t}", chunks[2] if len(chunks)>2 else '')
        await page.evaluate("(t)=>{document.getElementById('codeCol4').textContent=t}", chunks[3] if len(chunks)>3 else '')
        await page.wait_for_timeout(100)
    el = await page.query_selector('#canvas')
    await el.screenshot(path=str(out_path))

async def render_sections(page, template_path, info, base_out):
    pass


async def main():
    user_id = read_env('LUOGU_USER_ID', '322081')
    start_page = int(read_env('START_PAGE', '104'))
    end_page = int(read_env('END_PAGE', str(start_page)))
    output_dir = Path(read_env('OUTPUT_DIR', str(Path.cwd() / 'artifacts' / f'luogu-{user_id}' / 'images')))
    output_dir.mkdir(parents=True, exist_ok=True)
    cookie = read_env('LUOGU_COOKIE', '')
    template_path = Path(__file__).resolve().parents[1] / 'luogu-snap' / 'src' / 'template.html'
    async with async_playwright() as pw:
        browser = await pw.chromium.launch(headless=True)
        storage = Path(__file__).resolve().parents[1] / 'artifacts' / f'luogu-{user_id}' / 'storage_state.json'
        if storage.exists():
            context = await browser.new_context(storage_state=str(storage))
        else:
            context = await browser.new_context()
        headers = {'User-Agent': 'Mozilla/5.0'}
        if cookie:
            headers['Cookie'] = cookie
        await context.set_extra_http_headers(headers)
        page = await context.new_page()
        # 断点续跑状态
        state_path = output_dir / 'state.json'
        try:
            state = json.loads(state_path.read_text(encoding='utf-8'))
        except Exception:
            state = {'done': {}}
        ignore_state = os.getenv('IGNORE_STATE', '0') == '1'

        def write_gallery():
            files = sorted([f.name for f in output_dir.glob('*.png')])
            items = '\n'.join([f'<div class="item"><img src="{name}" /><div class="name">{name}</div></div>' for name in files])
            html = f'<!doctype html><html><head><meta charset="utf-8"/><style>body{{font-family:Consolas,monospace}} .grid{{display:grid;grid-template-columns:repeat(auto-fill,minmax(320px,1fr));gap:12px;padding:16px}} .item img{{width:100%;height:auto;border:1px solid #ddd;border-radius:6px}} .name{{font-size:12px;color:#666;margin-top:4px}}</style></head><body><div class="grid">{items}</div></body></html>'
            (output_dir / 'index.html').write_text(html, encoding='utf-8')

        sem = asyncio.Semaphore(int(os.getenv('MAX_CONCURRENCY', '6')))

        async def process_link(link: str):
            async with sem:
                p = await context.new_page()
                try:
                    info = await extract_detail(p, link)
                    rid = info.get('rid') or 'rid'
                    if state['done'].get(rid) and not ignore_state:
                        await p.close()
                        return
                    pid = info.get('problemId') or 'unknown'
                    stamp = (info.get('submitTime', '')).replace(' ', '').replace(':', '').replace('-', '').replace('.', '')[:20]
                    base = '_'.join([x for x in [rid, pid, stamp or ''] if x])
                    await p.goto(link)
                    try:
                        await p.wait_for_load_state('networkidle')
                    except Exception:
                        await p.wait_for_timeout(400)
                    try:
                        await p.wait_for_selector('.test-case', timeout=5000)
                        await p.wait_for_function("(() => { const c = document.querySelectorAll('.test-case').length; window.__tc__ = window.__tc__ || c; const stable = (window.__tc__ === c); window.__tc__ = c; return stable && c > 0; })()", timeout=2000)
                    except Exception:
                        pass
                    el_cases = await p.query_selector('section.full-container')
                    if not el_cases:
                        el_cases = await p.query_selector('main')
                    if el_cases:
                        out1 = output_dir / f'{base}_cases.png'
                        await el_cases.screenshot(path=str(out1))
                        print(f'Saved {out1}')
                        write_gallery()
                    try:
                        tab = p.locator('text=源代码').first
                        await tab.click()
                        try:
                            await p.wait_for_selector('pre, pre code', timeout=5000)
                            await p.wait_for_function("(() => { const el = document.querySelector('pre code') || document.querySelector('pre'); if (!el) return false; const html = el.innerHTML || ''; return html.includes('<span') || (el.textContent || '').length > 0; })()", timeout=2000)
                        except Exception:
                            await p.wait_for_timeout(400)
                    except Exception:
                        pass
                    el_code = await p.query_selector('pre, pre code')
                    if not el_code:
                        el_code = await p.query_selector('section.full-container')
                    if el_code:
                        out2 = output_dir / f'{base}_code.png'
                        await el_code.screenshot(path=str(out2))
                        print(f'Saved {out2}')
                        write_gallery()
                    state['done'][rid] = True
                    state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding='utf-8')
                except Exception as e:
                    print(f'Error processing {link}: {e}')
                finally:
                    await p.close()

        if end_page >= start_page:
            for pno in range(start_page, end_page + 1):
                links = await collect_links(page, user_id, pno)
                print(f'Page {pno}: {len(links)} records')
                tasks = [asyncio.create_task(process_link(link)) for link in links]
                await asyncio.gather(*tasks, return_exceptions=True)
        else:
            empty = 0
            pno = start_page
            while True:
                links = await collect_links(page, user_id, pno)
                print(f'Page {pno}: {len(links)} records')
                if not links:
                    empty += 1
                    if empty >= 2:
                        break
                    pno += 1
                    continue
                empty = 0
                tasks = [asyncio.create_task(process_link(link)) for link in links]
                await asyncio.gather(*tasks, return_exceptions=True)
                pno += 1
            write_gallery()
        await browser.close()
    print('Done. Images at', output_dir)


if __name__ == '__main__':
    asyncio.run(main())

